import json
import uuid
from typing import Dict, List, Optional, Any
from datetime import datetime, timedelta
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update, and_, or_
from sqlalchemy.orm import selectinload
from loguru import logger

from app.models.db import Task, User, UserSession, TaskConfig, ApiFunction, DataSet


class TaskService:
    """任务服务类 - 处理任务和任务结果的存储管理"""
    
    @staticmethod
    async def create_task(
        task_data: Dict[str, Any],
        user_id: str,
        session_id: Optional[str] = None,
        db: AsyncSession = None
    ) -> Task:
        """
        创建新任务
        
        Args:
            task_data: 任务数据
            user_id: 用户ID
            session_id: 会话ID（可选）
            db: 数据库会话
            
        Returns:
            创建的任务对象
        """
        try:
            task = Task(
                id=str(uuid.uuid4()),
                task_name=task_data.get("task_name", "未命名任务"),
                task_type=task_data.get("task_type", "unknown"),
                task_description=task_data.get("task_description"),
                priority=task_data.get("priority", 0),
                user_id=user_id,
                session_id=session_id,
                intent_id=task_data.get("intent_id"),
                parent_task_id=task_data.get("parent_task_id"),
                task_config_id=task_data.get("task_config_id"),
                api_function_id=task_data.get("api_function_id"),
                dataset_id=task_data.get("dataset_id"),
                generated_sql=task_data.get("generated_sql"),
                execution_order=task_data.get("execution_order", 1),
                estimated_duration=task_data.get("estimated_duration"),
                status="pending"
            )
            
            # 设置输入参数
            if task_data.get("input_parameters"):
                task.set_input_parameters(task_data["input_parameters"])
            
            # 设置任务配置
            if task_data.get("task_config"):
                task.set_task_config_json(task_data["task_config"])
            
            db.add(task)
            await db.commit()
            await db.refresh(task)
            
            logger.info(f"创建任务成功: {task.id} - {task.task_name}")
            return task
            
        except Exception as e:
            logger.error(f"创建任务失败: {str(e)}")
            await db.rollback()
            raise e
    
    @staticmethod
    async def create_task_from_intent(
        intent_result: Dict[str, Any],
        user_id: str,
        session_id: Optional[str] = None,
        db: AsyncSession = None
    ) -> Task:
        """
        从意图识别结果创建任务
        
        Args:
            intent_result: 意图识别结果
            user_id: 用户ID
            session_id: 会话ID（可选）
            db: 数据库会话
            
        Returns:
            创建的任务对象
        """
        try:
            task_data = {
                "task_name": intent_result.get("task_name", "AI任务"),
                "task_type": intent_result.get("intent_type", "unknown"),
                "task_description": intent_result.get("analysis_reason"),
                "intent_id": intent_result.get("intent_id"),
                "priority": intent_result.get("priority", 0),
                "execution_order": intent_result.get("execution_order", 1),
                "input_parameters": intent_result.get("extracted_params", {}),
                "task_config": {
                    "confidence": intent_result.get("confidence", 0.0),
                    "next_action": intent_result.get("next_action"),
                    "user_message": intent_result.get("user_message")
                }
            }
            
            # 根据意图类型设置特定字段
            if intent_result.get("intent_type") == "api_function":
                task_data["api_function_id"] = intent_result.get("api_function_id")
                task_data["task_config_id"] = intent_result.get("task_config_id")
            elif intent_result.get("intent_type") == "data_query":
                task_data["dataset_id"] = intent_result.get("dataset_id")
                task_data["generated_sql"] = intent_result.get("generated_sql")
            
            return await TaskService.create_task(task_data, user_id, session_id, db)
            
        except Exception as e:
            logger.error(f"从意图创建任务失败: {str(e)}")
            raise e
    
    @staticmethod
    async def get_task_by_id(task_id: str, db: AsyncSession) -> Optional[Task]:
        """
        根据ID获取任务
        
        Args:
            task_id: 任务ID
            db: 数据库会话
            
        Returns:
            任务对象或None
        """
        try:
            query = select(Task).where(Task.id == task_id)
            result = await db.execute(query)
            return result.scalar_one_or_none()
            
        except Exception as e:
            logger.error(f"获取任务失败: {task_id} - {str(e)}")
            return None
    
    @staticmethod
    async def get_tasks_by_session(
        session_id: str, 
        db: AsyncSession,
        limit: int = 50,
        offset: int = 0
    ) -> List[Task]:
        """
        根据会话ID获取任务列表
        
        Args:
            session_id: 会话ID
            db: 数据库会话
            limit: 限制数量
            offset: 偏移量
            
        Returns:
            任务列表
        """
        try:
            query = select(Task).where(
                Task.session_id == session_id
            ).order_by(Task.created_at.desc()).limit(limit).offset(offset)
            
            result = await db.execute(query)
            return result.scalars().all()
            
        except Exception as e:
            logger.error(f"获取会话任务失败: {session_id} - {str(e)}")
            return []
    
    @staticmethod
    async def get_tasks_by_user(
        user_id: str,
        db: AsyncSession,
        status_filter: Optional[List[str]] = None,
        limit: int = 50,
        offset: int = 0
    ) -> List[Task]:
        """
        获取用户的任务
        
        Args:
            user_id: 用户ID
            db: 数据库会话
            status_filter: 状态过滤器
            limit: 限制条数
            offset: 偏移量
            
        Returns:
            任务列表
        """
        try:
            query = select(Task).where(Task.user_id == user_id)
            
            if status_filter:
                query = query.where(Task.status.in_(status_filter))
            
            query = query.order_by(Task.created_at.desc()).limit(limit).offset(offset)
            
            result = await db.execute(query)
            return result.scalars().all()
            
        except Exception as e:
            logger.error(f"获取用户任务失败: {user_id} - {str(e)}")
            return []
    
    @staticmethod
    async def update_task_status(
        task_id: str,
        status: str,
        db: AsyncSession,
        progress: Optional[int] = None,
        error_message: Optional[str] = None
    ) -> bool:
        """
        更新任务状态
        
        Args:
            task_id: 任务ID
            status: 新状态
            db: 数据库会话
            progress: 进度百分比
            error_message: 错误信息
            
        Returns:
            是否更新成功
        """
        try:
            update_data = {
                "status": status,
                "updated_at": datetime.now()
            }
            
            if progress is not None:
                update_data["progress"] = progress
            
            if status == "running" and progress is None:
                update_data["started_at"] = datetime.now()
                update_data["progress"] = 0
            elif status in ["completed", "failed", "cancelled"]:
                update_data["completed_at"] = datetime.now()
                update_data["progress"] = 100 if status == "completed" else update_data.get("progress", 0)
            
            query = update(Task).where(Task.id == task_id).values(**update_data)
            result = await db.execute(query)
            await db.commit()
            
            return result.rowcount > 0
            
        except Exception as e:
            logger.error(f"更新任务状态失败: {task_id} - {str(e)}")
            await db.rollback()
            return False
    
    @staticmethod
    async def update_task_messageid(
        task_id: str,
        messageid: str,
        db: AsyncSession
    ) -> bool:
        """
        更新任务的messageid字段
        
        Args:
            task_id: 任务ID
            messageid: 系统回复消息ID
            db: 数据库会话
            
        Returns:
            是否更新成功
        """
        try:
            query = update(Task).where(Task.id == task_id).values(
                messageid=messageid,
                updated_at=datetime.now()
            )
            result = await db.execute(query)
            await db.commit()
            
            logger.info(f"更新任务messageid成功: task_id={task_id}, messageid={messageid}")
            return result.rowcount > 0
            
        except Exception as e:
            logger.error(f"更新任务messageid失败: task_id={task_id}, messageid={messageid} - {str(e)}")
            await db.rollback()
            return False
    
    @staticmethod
    async def update_tasks_messageid_by_session(
        session_id: str,
        messageid: str,
        db: AsyncSession,
        time_threshold_minutes: int = 5
    ) -> int:
        """
        批量更新会话中最近创建的任务的messageid字段
        
        Args:
            session_id: 会话ID
            messageid: 系统回复消息ID
            db: 数据库会话
            time_threshold_minutes: 时间阈值（分钟），只更新这个时间内创建的任务
            
        Returns:
            更新的任务数量
        """
        try:
            # 计算时间阈值
            time_threshold = datetime.now() - timedelta(minutes=time_threshold_minutes)
            
            # 更新最近创建的任务
            query = update(Task).where(
                and_(
                    Task.session_id == session_id,
                    Task.created_at >= time_threshold,
                    Task.messageid.is_(None)  # 只更新还没有设置messageid的任务
                )
            ).values(
                messageid=messageid,
                updated_at=datetime.now()
            )
            
            result = await db.execute(query)
            await db.commit()
            
            updated_count = result.rowcount
            logger.info(f"批量更新任务messageid成功: session_id={session_id}, messageid={messageid}, 更新数量={updated_count}")
            return updated_count
            
        except Exception as e:
            logger.error(f"批量更新任务messageid失败: session_id={session_id}, messageid={messageid} - {str(e)}")
            await db.rollback()
            return 0
    
    @staticmethod
    async def save_execution_result(
        task_id: str,
        execution_result: Dict[str, Any],
        db: AsyncSession
    ) -> Optional[Task]:
        """保存任务执行结果到Task表（增强数据量控制）"""
        try:
            from config import MAX_STORAGE_ROWS
            
            task = await TaskService.get_task_by_id(task_id, db)
            if not task:
                logger.error(f"任务不存在: {task_id}")
                return None

            # 计算执行时长
            started_at = execution_result.get("started_at", task.started_at or datetime.now())
            completed_at = execution_result.get("completed_at", datetime.now())
            
            if isinstance(started_at, str):
                started_at = datetime.fromisoformat(started_at.replace('Z', '+00:00'))
            if isinstance(completed_at, str):
                completed_at = datetime.fromisoformat(completed_at.replace('Z', '+00:00'))
            
            duration_ms = int((completed_at - started_at).total_seconds() * 1000)
            
            # 确定执行状态和结果类型
            execution_status = "success" if execution_result.get("status") == "success" else "error"
            result_type = execution_result.get("task_type", task.task_type)
            
            # 🎯 智能处理数据存储
            result_data = execution_result.get("data")
            storage_info = {
                "storage_truncated": False,
                "storage_message": None
            }
            
            if result_data and isinstance(result_data, list):
                original_count = len(result_data)
                
                if original_count > MAX_STORAGE_ROWS:
                    # 需要截断存储
                    result_data = result_data[:MAX_STORAGE_ROWS]
                    storage_info["storage_truncated"] = True
                    storage_info["storage_message"] = f"存储已截断：原始 {original_count} 条，存储前 {MAX_STORAGE_ROWS} 条"
                    logger.info(f"任务 {task_id} 数据存储已截断: {original_count} -> {MAX_STORAGE_ROWS}")
            
            # 更新任务的执行结果字段
            update_data = {
                "status": "completed" if execution_status == "success" else "failed",
                "progress": 100,
                "completed_at": completed_at,
                "updated_at": datetime.now(),
                "execution_status": execution_status,
                "result_type": result_type,
                "duration_ms": duration_ms,
                "result_summary": execution_result.get("message"),
                "error_message": execution_result.get("message") if execution_status == "error" else None,
                "execution_environment": "python_async",
                "execution_node": "local"
            }
            
            # 设置结果数据（可能已截断）
            if result_data:
                task.set_result_data(result_data)
            
            # 🎯 记录截断和数据量信息到资源使用字段
            metadata = {
                "query_truncated": execution_result.get("is_truncated", False),
                "total_rows_returned": execution_result.get("total_rows_returned"),
                "truncation_info": execution_result.get("truncation_info"),
                "display_strategy": execution_result.get("display_strategy"),
                **storage_info
            }
            task.set_resource_usage(metadata)
            
            # 根据任务类型设置特定字段
            if result_type == "data_query":
                update_data["sql_query"] = execution_result.get("sql_query")
                update_data["rows_affected"] = execution_result.get("total_rows_returned")
                task.set_columns_info(execution_result.get("columns", []))
            elif result_type == "api_function":
                update_data["api_endpoint"] = execution_result.get("function_name")
            
            # 执行更新
            query = update(Task).where(Task.id == task_id).values(**update_data)
            await db.execute(query)
            await db.commit()
            await db.refresh(task)
            
            logger.info(f"保存任务执行结果成功: {task_id}")
            return task
            
        except Exception as e:
            logger.error(f"保存任务执行结果失败: {task_id} - {str(e)}")
            await db.rollback()
            raise e
    
    @staticmethod
    async def get_execution_results(
        task_id: str,
        db: AsyncSession
    ) -> List[Task]:
        """
        获取任务的执行结果（兼容方法，实际返回任务本身）
        
        Args:
            task_id: 任务ID
            db: 数据库会话
            
        Returns:
            任务列表（包含执行结果）
        """
        try:
            task = await TaskService.get_task_by_id(task_id, db)
            if task:
                return [task]  # 返回包含任务的列表以保持兼容性
            return []
            
        except Exception as e:
            logger.error(f"获取任务执行结果失败: {task_id} - {str(e)}")
            return []
    
    @staticmethod
    async def get_latest_execution_result(
        task_id: str,
        db: AsyncSession
    ) -> Optional[Task]:
        """
        获取任务的最新执行结果（兼容方法，实际返回任务本身）
        
        Args:
            task_id: 任务ID
            db: 数据库会话
            
        Returns:
            任务对象（包含执行结果）
        """
        try:
            return await TaskService.get_task_by_id(task_id, db)
            
        except Exception as e:
            logger.error(f"获取最新任务执行结果失败: {task_id} - {str(e)}")
            return None
    
    @staticmethod
    async def retry_failed_task(
        task_id: str,
        db: AsyncSession
    ) -> bool:
        """
        重试失败的任务（已废弃，因为不支持重试机制）
        
        Args:
            task_id: 任务ID
            db: 数据库会话
            
        Returns:
            是否可以重试（始终返回False）
        """
        logger.warning(f"重试功能已废弃，任务ID: {task_id}")
        return False
    
    @staticmethod
    async def cancel_task(
        task_id: str,
        db: AsyncSession
    ) -> bool:
        """
        取消任务
        
        Args:
            task_id: 任务ID
            db: 数据库会话
            
        Returns:
            是否取消成功
        """
        try:
            return await TaskService.update_task_status(task_id, "cancelled", db)
            
        except Exception as e:
            logger.error(f"取消任务失败: {task_id} - {str(e)}")
            return False
    
    @staticmethod
    async def get_pending_tasks(
        db: AsyncSession,
        limit: int = 10
    ) -> List[Task]:
        """
        获取待执行的任务
        
        Args:
            db: 数据库会话
            limit: 限制条数
            
        Returns:
            待执行任务列表
        """
        try:
            query = select(Task).where(
                Task.status == "pending"
            ).order_by(
                Task.priority.desc(),
                Task.execution_order.asc(),
                Task.created_at.asc()
            ).limit(limit)
            
            result = await db.execute(query)
            return result.scalars().all()
            
        except Exception as e:
            logger.error(f"获取待执行任务失败: {str(e)}")
            return []
    
    @staticmethod
    async def get_task_statistics(
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        db: AsyncSession = None
    ) -> Dict[str, Any]:
        """
        获取任务统计信息
        
        Args:
            user_id: 用户ID（可选）
            session_id: 会话ID（可选）
            db: 数据库会话
            
        Returns:
            统计信息
        """
        try:
            base_query = select(Task)
            
            if user_id:
                base_query = base_query.where(Task.user_id == user_id)
            if session_id:
                base_query = base_query.where(Task.session_id == session_id)
            
            # 获取各状态的任务数量
            from sqlalchemy import func
            
            status_query = select(
                Task.status,
                func.count(Task.id).label('count')
            ).group_by(Task.status)
            
            if user_id:
                status_query = status_query.where(Task.user_id == user_id)
            if session_id:
                status_query = status_query.where(Task.session_id == session_id)
            
            result = await db.execute(status_query)
            status_counts = {row[0]: row[1] for row in result.fetchall()}
            
            # 计算总数
            total_tasks = sum(status_counts.values())
            
            return {
                "total_tasks": total_tasks,
                "status_counts": status_counts,
                "success_rate": (status_counts.get("completed", 0) / total_tasks * 100) if total_tasks > 0 else 0
            }
            
        except Exception as e:
            logger.error(f"获取任务统计失败: {str(e)}")
            return {
                "total_tasks": 0,
                "status_counts": {},
                "success_rate": 0
            } 